import os
import pandas as pd
import PharseSpliter
import re

# 自定义排序函数，提取数字进行比较
def sort_key(file_name):
    # 使用正则表达式找到数字部分
    numbers = re.findall(r'\d+', file_name)
    if numbers:
        # 返回第一个找到的数字，作为排序依据
        return int(numbers[0])
    return 0  # 如果没有找到数字，返回0

# 音频处理函数
def function(audio_url, text, save_path):
    print(audio_url, text, save_path)
    pass  # 这里是你的音频处理代码

# 音频文件总文件夹
sph = "/root/autodl-tmp/TEDspliter/train/sph"
# 字幕文件总文件夹
stm = "/root/autodl-tmp/TEDspliter/train/stm_without_special_tokens"
# 输出文件夹地址
out_path = "/root/autodl-tmp/TEDspliter/increment_TED/train"

# 遍历音频文件总文件夹
for audio_folder in os.listdir(sph):
    # 音频文件夹的绝对路径
    audio_folder_path = os.path.join(sph, audio_folder)
    # 对应的字幕文件的绝对路径
    subtitle_file_path = os.path.join(stm, f"{audio_folder}.csv")

    # 读取字幕文件
    df = pd.read_csv(subtitle_file_path)
    subtitles = df['caption'].tolist()

    # 组装保存地址
    audio_save_path = os.path.join(out_path,'sph',audio_folder)
    csv_save_path = os.path.join(out_path,'stm',audio_folder)
    
    # 遍历音频文件夹
    for i, audio_file in enumerate(sorted(os.listdir(audio_folder_path), key=sort_key)):
        # 音频文件的绝对路径
        audio_file_path = os.path.join(audio_folder_path, audio_file)
        # 对应的字幕
        text = subtitles[i]
        
        # 组装保存地址
        # 获取音频文件名(segement_n.wav)
        basename = os.path.basename(audio_file)
        basename = os.path.splitext(basename)[0]

        _audio_save_path = os.path.join(audio_save_path,basename)
        _csv_save_path = os.path.join(csv_save_path,basename+'.csv')

        # 调用音频处理函数
        PharseSpliter.spliter_to_pharse(audio_url=audio_file_path,
                                text=text,
                                save_path_sph=_audio_save_path,
                                save_path_stm=_csv_save_path)

print("全部处理完毕")